"""
主要颜色检测器 - 基于K-Means聚类的屏幕中心区域主要颜色检测
参考 test5.py 的实现思路，专门用于游戏颜色触发
"""

import numpy as np
import pyautogui
from PIL import ImageGrab
import time
import logging
import warnings

logger = logging.getLogger(__name__)

class DominantColorDetector:
    """主要颜色检测器 - 使用K-Means提取屏幕区域的主要颜色"""
    
    def __init__(self, region_size=200, k_clusters=3, cache_duration=0.1):
        """
        初始化主要颜色检测器
        
        Args:
            region_size: 检测区域边长（像素）
            k_clusters: K-Means聚类数量
            cache_duration: 缓存持续时间（秒）
        """
        self.region_size = region_size
        self.k_clusters = k_clusters
        self.cache_duration = cache_duration
        
        # 缓存机制
        self.cache = {}
        self.last_detection_time = 0
        self.last_dominant_color = None
        
        # 性能统计
        self.detection_times = []
        self.total_detections = 0
        
    def get_screen_center_region(self, center_x=None, center_y=None, region_size=None):
        """
        获取屏幕中心区域的截图
        
        Args:
            center_x: 中心点X坐标（默认屏幕中心）
            center_y: 中心点Y坐标（默认屏幕中心）
            region_size: 区域大小（默认使用初始化值）
            
        Returns:
            PIL Image对象或None
        """
        try:
            # 使用默认值
            if region_size is None:
                region_size = self.region_size
                
            # 获取屏幕尺寸
            screen_width, screen_height = pyautogui.size()
            
            # 计算中心点
            if center_x is None:
                center_x = screen_width // 2
            if center_y is None:
                center_y = screen_height // 2
            
            # 计算区域边界
            half_size = region_size // 2
            left = max(0, center_x - half_size)
            top = max(0, center_y - half_size)
            right = min(screen_width, center_x + half_size)
            bottom = min(screen_height, center_y + half_size)
            
            # 截取区域
            bbox = (left, top, right, bottom)
            screenshot = ImageGrab.grab(bbox=bbox)
            
            logger.debug(f"截取区域: {bbox}, 大小: {screenshot.size}")
            return screenshot
            
        except Exception as e:
            logger.error(f"获取屏幕区域失败: {e}")
            return None
    
    def extract_dominant_color(self, image, k_clusters=None):
        """
        使用K-Means提取图像的主要颜色
        
        Args:
            image: PIL Image对象
            k_clusters: 聚类数量（默认使用初始化值）
            
        Returns:
            tuple: (主要颜色RGB, 所有聚类中心, 权重)
        """
        try:
            start_time = time.time()
            
            if k_clusters is None:
                k_clusters = self.k_clusters
            
            # 转换为numpy数组
            img_array = np.array(image)
            
            # 重塑为像素数组 (n_pixels, 3)
            pixels = img_array.reshape(-1, 3)
            
            # 检查唯一颜色数量
            unique_colors = np.unique(pixels, axis=0)
            actual_clusters = min(k_clusters, len(unique_colors))
            
            # 如果只有一种颜色，直接返回
            if len(unique_colors) == 1:
                dominant_color = unique_colors[0]
                detection_time = time.time() - start_time
                self.detection_times.append(detection_time)
                logger.debug(f"单一颜色检测: {dominant_color}, 耗时: {detection_time:.3f}s")
                return dominant_color, unique_colors, np.array([1.0])
            
            # 执行K-Means聚类，抑制警告
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore", category=UserWarning)
                warnings.filterwarnings("ignore", message=".*Number of distinct clusters.*")
                try:
                    from sklearn.exceptions import ConvergenceWarning
                    warnings.filterwarnings("ignore", category=ConvergenceWarning)
                except ImportError:
                    pass
                
                from sklearn.cluster import KMeans
                kmeans = KMeans(n_clusters=actual_clusters, n_init=10, random_state=42)
                kmeans.fit(pixels)
            
            # 获取聚类中心（主要颜色）
            cluster_centers = kmeans.cluster_centers_.astype(int)
            
            # 计算每个聚类的权重
            labels = kmeans.labels_
            unique_labels, counts = np.unique(labels, return_counts=True)
            
            # 找到最大聚类（主要颜色）
            dominant_cluster_idx = np.argmax(counts)
            dominant_color = cluster_centers[dominant_cluster_idx]
            
            # 计算权重
            weights = counts / len(pixels)
            
            detection_time = time.time() - start_time
            self.detection_times.append(detection_time)
            self.total_detections += 1
            
            logger.debug(f"K-Means检测完成: 主要颜色={dominant_color}, 聚类数={len(cluster_centers)}, 耗时={detection_time:.3f}s")
            
            return dominant_color, cluster_centers, weights
            
        except ImportError:
            logger.error("scikit-learn未安装，无法使用K-Means检测")
            # 回退到简单平均
            img_array = np.array(image)
            pixels = img_array.reshape(-1, 3)
            avg_color = np.mean(pixels, axis=0).astype(int)
            return avg_color, np.array([avg_color]), np.array([1.0])
        except Exception as e:
            logger.error(f"K-Means颜色提取失败: {e}")
            return None, None, None
    
    def get_dominant_color_at_center(self, center_x=None, center_y=None, region_size=None, use_cache=True):
        """
        获取屏幕中心区域的主要颜色
        
        Args:
            center_x: 中心点X坐标
            center_y: 中心点Y坐标  
            region_size: 检测区域大小
            use_cache: 是否使用缓存
            
        Returns:
            tuple: (主要颜色RGB, 置信度, 检测信息)
        """
        try:
            current_time = time.time()
            
            # 检查缓存
            if use_cache and self.last_dominant_color is not None:
                if current_time - self.last_detection_time < self.cache_duration:
                    logger.debug("使用缓存的主要颜色")
                    return self.last_dominant_color, 1.0, {"cached": True}
            
            # 获取屏幕区域
            screenshot = self.get_screen_center_region(center_x, center_y, region_size)
            if screenshot is None:
                return None, 0.0, {"error": "截图失败"}
            
            # 提取主要颜色
            dominant_color, all_colors, weights = self.extract_dominant_color(screenshot)
            if dominant_color is None:
                return None, 0.0, {"error": "颜色提取失败"}
            
            # 计算置信度（主要颜色的权重）
            confidence = weights[0] if len(weights) > 0 else 1.0
            
            # 更新缓存
            self.last_dominant_color = tuple(dominant_color)
            self.last_detection_time = current_time
            
            # 检测信息
            detection_info = {
                "dominant_color": tuple(dominant_color),
                "all_colors": [tuple(color) for color in all_colors],
                "weights": weights.tolist(),
                "confidence": confidence,
                "region_size": screenshot.size,
                "detection_time": self.detection_times[-1] if self.detection_times else 0
            }
            
            return tuple(dominant_color), confidence, detection_info
            
        except Exception as e:
            logger.error(f"获取主要颜色失败: {e}")
            return None, 0.0, {"error": str(e)}
    
    def is_color_match_dominant(self, target_color, tolerance=30, center_x=None, center_y=None, region_size=None):
        """
        检查屏幕中心区域的主要颜色是否匹配目标颜色
        
        Args:
            target_color: 目标颜色 (R, G, B)
            tolerance: 颜色容差
            center_x: 检测中心X坐标
            center_y: 检测中心Y坐标
            region_size: 检测区域大小
            
        Returns:
            tuple: (是否匹配, 置信度, 检测信息)
        """
        try:
            # 获取主要颜色
            dominant_color, color_confidence, detection_info = self.get_dominant_color_at_center(
                center_x, center_y, region_size
            )
            
            if dominant_color is None:
                return False, 0.0, detection_info
            
            # 计算颜色距离
            target_rgb = np.array(target_color)
            dominant_rgb = np.array(dominant_color)
            color_distance = np.linalg.norm(target_rgb - dominant_rgb)
            
            # 判断是否匹配
            is_match = color_distance <= tolerance
            
            # 计算综合置信度（颜色置信度 × 匹配度）
            if is_match:
                match_confidence = max(0, 1 - color_distance / tolerance)
                final_confidence = color_confidence * match_confidence
            else:
                final_confidence = 0.0
            
            # 更新检测信息
            detection_info.update({
                "target_color": target_color,
                "color_distance": color_distance,
                "tolerance": tolerance,
                "is_match": is_match,
                "match_confidence": final_confidence
            })
            
            logger.debug(f"颜色匹配检测: 目标={target_color}, 主要={dominant_color}, 距离={color_distance:.1f}, 匹配={is_match}")
            
            return is_match, final_confidence, detection_info
            
        except Exception as e:
            logger.error(f"主要颜色匹配检测失败: {e}")
            return False, 0.0, {"error": str(e)}
    
    def get_performance_stats(self):
        """获取性能统计信息"""
        if not self.detection_times:
            return {
                "total_detections": 0,
                "avg_detection_time": 0,
                "min_detection_time": 0,
                "max_detection_time": 0
            }
        
        return {
            "total_detections": self.total_detections,
            "avg_detection_time": np.mean(self.detection_times),
            "min_detection_time": np.min(self.detection_times),
            "max_detection_time": np.max(self.detection_times),
            "cache_hit_rate": 1 - len(self.detection_times) / max(1, self.total_detections)
        }
    
    def clear_cache(self):
        """清除缓存"""
        self.cache.clear()
        self.last_dominant_color = None
        self.last_detection_time = 0
        logger.debug("缓存已清除")

# 便捷函数
def get_screen_dominant_color(region_size=200, k_clusters=3, center_x=None, center_y=None):
    """
    便捷函数：获取屏幕中心区域的主要颜色
    
    Returns:
        tuple: (主要颜色RGB, 置信度)
    """
    detector = DominantColorDetector(region_size, k_clusters)
    dominant_color, confidence, _ = detector.get_dominant_color_at_center(center_x, center_y)
    return dominant_color, confidence

def check_dominant_color_match(target_color, tolerance=30, region_size=200, k_clusters=3):
    """
    便捷函数：检查屏幕中心主要颜色是否匹配目标
    
    Returns:
        bool: 是否匹配
    """
    detector = DominantColorDetector(region_size, k_clusters)
    is_match, _, _ = detector.is_color_match_dominant(target_color, tolerance)
    return is_match
